from tree_sitter import Language, Parser, Tree, Node
import tree_sitter_cpp
import os, json


class AST:
    def __init__(self, code: bytes):
        self.tree: Tree = self.parse(code)
        self.insert_rules: dict = {}
        self.split_rules: dict = {}
        self.insert_points_list: list = []
        self.sub_tree_list: list = []
        self.text: bytes = self.tree.root_node.text

    def parse(self, code: bytes):
        CPP_LANGUATE = Language(tree_sitter_cpp.language())
        parser = Parser(CPP_LANGUATE)
        try:
            return parser.parse(code)
        except Exception as e:
            print(f"Failed to parse ast with exception: {e}")
            return
        
    def compile(self):
        def recover_code(node: Node) -> bytes:
            if len(node.children) != 0:
                code_fragments = []
                for children in node.children:
                    code_fragments.append(recover_code(children))
                code = b' '.join(code_fragments)
                return code
            else:
                return node.text

    def insert_point_get(self) -> list:
        def get_points(node: Node) -> list:
            insert_points = []
            if node.type in self.insert_rules:
                insert_points.append(node)
            else:
                pass
            for child in node.named_children:
                if insert_points:
                    insert_points.extend(get_points(child))
                else:
                    insert_points = get_points(child)
            return insert_points

        def classification(node: Node) -> list:
            if node.type in self.insert_rules:
                for rule in self.insert_rules[node.type]:
                    if all(eval(condition,{},{'node':node}) for condition in rule['condition'].values() if condition):
                        return rule['insert_position']
                    else:
                        return []
                        
        try:
            with open(f'{os.path.dirname(os.path.abspath(__file__))}/rules/insert_rules.json', 'r') as f:
                self.insert_rules = json.load(f)
        except Exception as e:
            print(f"Failed to load insert rule with exception: {e}")
        insert_points = get_points(self.tree.root_node)
        for insert_point in insert_points:
            insert_position = classification(insert_point)
            if insert_position:
                self.insert_points_list.append({'node': insert_point, 'insert_position': insert_position})
        return
    
    def code_split(self) -> list:
        def split(node: Node) -> list:
            split_points = []
            if node.type in self.split_rules and judge(node):
                split_points.append(node)
            else:
                for child in node.named_children:
                    if split_points:
                        split_points.extend(split(child))
                    else:
                        split_points = split(child)
            return split_points
        
        def judge(node: Node) -> list:
            for rule in self.split_rules[node.type]:
                if all(eval(condition,{},{'node':node}) for condition in rule['condition'].values() if condition):
                    return True
                else:
                    return False
            
        try:
            with open(f'{os.path.dirname(os.path.abspath(__file__))}/rules/split_rules.json', 'r') as f:
                self.split_rules = json.load(f)
        except Exception as e:
            print(f"Failed to load split rule with exception: {e}")
        self.sub_tree_list = split(self.tree.root_node)
        return

    def code_merge(self, another_tree: 'AST') -> bool:
        sub_trees = another_tree.sub_tree_list
        offset = 0
        for insert_point in self.insert_points_list:
            while sub_trees:
                sub_trees_len = len(sub_trees)
                for insert_position in insert_point['insert_position']:
                    if sub_trees[0].type in insert_position['nodes']:
                        insert_pos = 0
                        if insert_position['position'] == 'prev_sibling':
                            insert_pos = insert_point['node'].start_byte + offset
                        elif insert_position['position'] == 'next_sibling':
                            insert_pos = insert_point['node'].end_byte + offset
                        elif insert_position['position'] == 'left_child':
                            insert_pos = insert_point['node'].named_children[0].start_byte + offset

                        filler = b''
                        if sub_trees[0].text.endswith(b'}'):
                            filler = b';\r\n'
                        elif sub_trees[0].text.endswith(b';'):
                            filler = b'\r\n'

                        new_text = self.text[:insert_pos] + sub_trees[0].text + filler + self.text[insert_pos:]
                        self.text = new_text
                        offset += sub_trees[0].end_byte - sub_trees[0].start_byte + len(filler)
                        sub_trees.remove(sub_trees[0])
                        break
                    else:
                        pass
                if sub_trees_len == len(sub_trees):
                    break
        if sub_trees:
            return False
        else:
            return True
    
    def code_output(self, path: str) -> None:
        with open(path, 'wb') as f:
            f.write(self.text)
        return